use ndarray::prelude::*;
use ndarray::{Array, Array1, Array2};
use ndarray_linalg::Solve;
use std::collections::VecDeque;
use std::f64::consts::PI;

// 数据结构：存储时间序列数据点
#[derive(Debug, Clone)]
struct DataPoint {
    timestamp: f64, // 时间戳（如月份：1.0, 2.0, ...）
    value: f64,     // 该时间点的值（如报废数量）
}

// 趋势+季节性模型
struct TrendSeasonalModel {
    trend_coefficients: Array1<f64>,    // [截距, 斜率]
    seasonal_coefficients: Array1<f64>, // 傅里叶系数
    period: f64,                        // 周期长度
    fourier_terms: usize,               // 傅里叶项数
    last_train_time: f64,               // 最后训练时间
}

// 预测系统配置
struct ForecastConfig {
    max_history: usize,    // 最大历史数据量（月）
    min_history: usize,    // 最小历史数据量（月）
    retrain_interval: f64, // 重新训练间隔（月）
    data_decay_rate: f64,  // 数据衰减率
    period: f64,           // 季节性周期（月）
    fourier_terms: usize,  // 傅里叶项数
}

// 预测系统状态
struct ForecastingSystem {
    config: ForecastConfig,
    data_buffer: VecDeque<DataPoint>, // 数据缓冲区（环形）
    model: Option<TrendSeasonalModel>, // 当前模型
    last_retrain_time: f64,           // 最后训练时间
}

impl TrendSeasonalModel {
    fn new() -> Self {
        Self {
            trend_coefficients: Array1::zeros(2),
            seasonal_coefficients: Array1::zeros(0),
            period: 12.0,
            fourier_terms: 2,
            last_train_time: 0.0,
        }
    }
    
    // 使用加权最小二乘法拟合模型
    fn fit_weighted(
        times: &Array1<f64>, 
        values: &Array1<f64>, 
        weights: &Array1<f64>,
        period: f64,
        fourier_terms: usize,
    ) -> Option<Self> {
        let n = times.len();
        if n < 12 {
            return None; // 数据不足
        }
        
        let total_terms = 2 + 2 * fourier_terms; // 趋势(2) + 傅里叶(2*terms)
        let mut features = Array2::zeros((n, total_terms));
        let mut weighted_values = Array1::zeros(n);
        
        // 构建特征矩阵和加权值向量
        for i in 0..n {
            let t = times[i];
            let w = weights[i].sqrt(); // 使用平方根权重进行加权最小二乘
            
            // 趋势特征
            features[[i, 0]] = w * 1.0;      // 常数项
            features[[i, 1]] = w * t;        // 时间项
            
            // 季节性特征（傅里叶级数）
            for k in 1..=fourier_terms {
                let angle = 2.0 * PI * k as f64 * t / period;
                features[[i, 2 + 2*k - 2]] = w * angle.sin(); // sin项
                features[[i, 2 + 2*k - 1]] = w * angle.cos(); // cos项
            }
            
            weighted_values[i] = w * values[i];
        }
        
        // 求解加权最小二乘问题
        let xt = features.t();
        let xtx = xt.dot(&features);
        let xty = xt.dot(&weighted_values);
        
        match xtx.solve(&xty) {
            Ok(all_coefficients) => {
                let trend_coefficients = all_coefficients.slice(s![0..2]).to_owned();
                let seasonal_coefficients = all_coefficients.slice(s![2..]).to_owned();
                
                Some(Self {
                    trend_coefficients,
                    seasonal_coefficients,
                    period,
                    fourier_terms,
                    last_train_time: *times.iter().max_by(|a, b| a.partial_cmp(b).unwrap()).unwrap_or(&0.0),
                })
            },
            Err(_) => None,
        }
    }
    
    // 预测指定时间点的值
    fn predict(&self, time: f64) -> f64 {
        // 趋势成分
        let trend = self.trend_coefficients[0] + self.trend_coefficients[1] * time;
        
        // 季节性成分
        let mut seasonal = 0.0;
        for k in 0..self.fourier_terms {
            let angle = 2.0 * PI * (k + 1) as f64 * time / self.period;
            seasonal += self.seasonal_coefficients[2*k] * angle.sin() 
                     + self.seasonal_coefficients[2*k + 1] * angle.cos();
        }
        
        trend + seasonal
    }
    
    // 分解预测值为趋势和季节性成分
    fn decompose(&self, time: f64) -> (f64, f64) {
        let trend = self.trend_coefficients[0] + self.trend_coefficients[1] * time;
        
        let mut seasonal = 0.0;
        for k in 0..self.fourier_terms {
            let angle = 2.0 * PI * (k + 1) as f64 * time / self.period;
            seasonal += self.seasonal_coefficients[2*k] * angle.sin() 
                     + self.seasonal_coefficients[2*k + 1] * angle.cos();
        }
        
        (trend, seasonal)
    }
}

impl ForecastConfig {
    fn default() -> Self {
        Self {
            max_history: 36,   // 最多保留3年数据
            min_history: 12,   // 至少需要1年数据
            retrain_interval: 1.0, // 每月重新训练
            data_decay_rate: 0.2,  // 每月衰减20%
            period: 12.0,      // 年度周期
            fourier_terms: 2,  // 使用2个傅里叶项
        }
    }
}

impl ForecastingSystem {
    fn new(config: ForecastConfig) -> Self {
        Self {
            config,
            data_buffer: VecDeque::with_capacity(36),
            model: None,
            last_retrain_time: 0.0,
        }
    }
    
    // 添加新数据点
    fn add_data_point(&mut self, timestamp: f64, value: f64) {
        // 维护数据缓冲区大小
        if self.data_buffer.len() >= self.config.max_history {
            self.data_buffer.pop_front();
        }
        
        self.data_buffer.push_back(DataPoint { timestamp, value });
        
        // 检查是否需要重新训练
        if self.should_retrain(timestamp) {
            self.retrain(timestamp);
        }
    }
    
    // 判断是否需要重新训练
    fn should_retrain(&self, current_time: f64) -> bool {
        if self.data_buffer.len() < self.config.min_history {
            return false; // 数据不足
        }
        
        // 检查是否到达重新训练时间
        current_time - self.last_retrain_time >= self.config.retrain_interval
    }
    
    // 重新训练模型
    fn retrain(&mut self, current_time: f64) {
        let n = self.data_buffer.len();
        
        // 准备数据数组
        let mut times = Array1::zeros(n);
        let mut values = Array1::zeros(n);
        let mut weights = Array1::zeros(n);
        
        for (i, point) in self.data_buffer.iter().enumerate() {
            times[i] = point.timestamp;
            values[i] = point.value;
            
            // 计算数据权重（近期数据权重高）
            let months_ago = current_time - point.timestamp;
            weights[i] = (-self.config.data_decay_rate * months_ago).exp();
        }
        
        // 训练新模型
        self.model = TrendSeasonalModel::fit_weighted(
            &times,
            &values,
            &weights,
            self.config.period,
            self.config.fourier_terms,
        );
        
        self.last_retrain_time = current_time;
    }
    
    // 预测未来值
    fn predict(&self, future_time: f64) -> Option<f64> {
        self.model.as_ref().map(|model| model.predict(future_time))
    }
    
    // 预测多个未来时间点
    fn predict_range(&self, start_time: f64, end_time: f64, step: f64) -> Vec<(f64, f64)> {
        let mut predictions = Vec::new();
        let mut current_time = start_time;
        
        while current_time <= end_time {
            if let Some(prediction) = self.predict(current_time) {
                predictions.push((current_time, prediction));
            }
            current_time += step;
        }
        
        predictions
    }
    
    // 计算库存使用时间
    fn estimate_inventory_duration(&self, current_inventory: f64, future_months: usize) -> Option<f64> {
        if self.data_buffer.is_empty() {
            return None;
        }
        
        let last_time = self.data_buffer.back().unwrap().timestamp;
        let future_times: Vec<f64> = (1..=future_months).map(|i| last_time + i as f64).collect();
        
        // 预测未来几个月的报废量
        let predictions: Vec<f64> = future_times.iter()
            .filter_map(|&t| self.predict(t))
            .collect();
        
        if predictions.is_empty() {
            return None;
        }
        
        // 计算平均月报废量
        let avg_monthly = predictions.iter().sum::<f64>() / predictions.len() as f64;
        
        if avg_monthly <= 0.0 {
            return Some(f64::INFINITY);
        }
        
        Some(current_inventory / avg_monthly)
    }
    
    // 获取模型性能指标
    fn model_performance(&self) -> Option<f64> {
        if self.data_buffer.len() < 6 || self.model.is_none() {
            return None;
        }
        
        let model = self.model.as_ref().unwrap();
        let n = self.data_buffer.len();
        
        // 使用最近6个月数据计算平均绝对百分比误差
        let start_idx = n.saturating_sub(6);
        let mut total_error = 0.0;
        let mut count = 0;
        
        for i in start_idx..n {
            let point = &self.data_buffer[i];
            let prediction = model.predict(point.timestamp);
            let error = (prediction - point.value).abs() / point.value.max(1.0);
            total_error += error;
            count += 1;
        }
        
        if count > 0 {
            Some(total_error / count as f64)
        } else {
            None
        }
    }
}

// 示例使用
fn main() {
    // 创建预测系统
    let config = ForecastConfig::default();
    let mut system = ForecastingSystem::new(config);
    
    // 模拟历史数据（包含趋势和季节性）
    let mut current_time = 1.0;
    for month in 1..=36 {
        // 模拟趋势：前24个月增长，后12个月下降
        let trend = if month <= 24 {
            100.0 + 2.0 * month as f64
        } else {
            148.0 - 1.0 * (month as f64 - 24.0)
        };
        
        // 模拟季节性：年度周期，振幅20
        let seasonal = 20.0 * (2.0 * PI * month as f64 / 12.0).sin();
        
        // 添加随机噪声
        let noise = rand::random::<f64>() * 10.0 - 5.0;
        
        let value = trend + seasonal + noise;
        
        system.add_data_point(current_time, value);
        current_time += 1.0;
    }
    
    // 预测未来
    let last_time = system.data_buffer.back().unwrap().timestamp;
    let predictions = system.predict_range(last_time + 1.0, last_time + 12.0, 1.0);
    
    println!("未来12个月预测:");
    for (month, prediction) in predictions {
        println!("  第{:.0}个月: {:.1}", month, prediction);
    }
    
    // 计算库存使用时间
    let current_inventory = 5000.0;
    if let Some(months_remaining) = system.estimate_inventory_duration(current_inventory, 12) {
        println!("当前库存 {} 预计还能使用 {:.1} 个月", current_inventory, months_remaining);
    }
    
    // 检查模型性能
    if let Some(mape) = system.model_performance() {
        println!("模型平均绝对百分比误差: {:.2}%", mape * 100.0);
    }
    
    // 添加新数据点并更新模型
    println!("\n添加新数据点并更新模型...");
    let new_value = 130.0; // 模拟新数据
    system.add_data_point(current_time, new_value);
    
    // 检查更新后的预测
    if let Some(new_prediction) = system.predict(current_time + 1.0) {
        println!("更新后下个月预测: {:.1}", new_prediction);
    }
}

// 辅助函数：生成测试数据
fn generate_test_data() -> Vec<DataPoint> {
    let mut data = Vec::new();
    let mut current_time = 1.0;
    
    for month in 1..=36 {
        // 生成包含趋势和季节性的测试数据
        let trend = 100.0 + 1.5 * month as f64;
        let seasonal = 15.0 * (2.0 * PI * month as f64 / 12.0).sin();
        let noise = rand::random::<f64>() * 8.0 - 4.0;
        
        data.push(DataPoint {
            timestamp: current_time,
            value: trend + seasonal + noise,
        });
        
        current_time += 1.0;
    }
    
    data
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_trend_seasonal_model() {
        let test_data = generate_test_data();
        let n = test_data.len();
        
        let times: Array1<f64> = Array1::from_iter(test_data.iter().map(|p| p.timestamp));
        let values: Array1<f64> = Array1::from_iter(test_data.iter().map(|p| p.value));
        let weights = Array1::ones(n);
        
        let model = TrendSeasonalModel::fit_weighted(
            &times, &values, &weights, 12.0, 2
        ).unwrap();
        
        // 测试预测功能
        let prediction = model.predict(37.0);
        assert!(prediction > 0.0);
        
        // 测试分解功能
        let (trend, seasonal) = model.decompose(37.0);
        assert!((prediction - (trend + seasonal)).abs() < 1e-10);
    }
    
    #[test]
    fn test_forecasting_system() {
        let config = ForecastConfig::default();
        let mut system = ForecastingSystem::new(config);
        
        // 添加测试数据
        let test_data = generate_test_data();
        for point in test_data {
            system.add_data_point(point.timestamp, point.value);
        }
        
        // 测试预测
        let last_time = system.data_buffer.back().unwrap().timestamp;
        let prediction = system.predict(last_time + 1.0);
        assert!(prediction.is_some());
        
        // 测试库存估算
        let duration = system.estimate_inventory_duration(1000.0, 12);
        assert!(duration.is_some());
        assert!(duration.unwrap() > 0.0);
    }
}